import pyaudio
import numpy as np
import asyncio
import time
from concurrent.futures import ThreadPoolExecutor

# 语音播放模块：处理音频播放

# 默认音频配置
DEFAULT_FORMAT = pyaudio.paInt16
DEFAULT_CHANNELS = 1
DEFAULT_RATE = 24000  # 修改为24kHz以匹配服务器
CHUNK_SIZE = 1440  # 24kHz * 60ms = 1440采样点

class AudioPlayer:
    def __init__(self, audio_queue):
        self.audio_queue = audio_queue
        self.is_playing = False
        # 创建线程池用于执行阻塞的音频操作
        self.executor = ThreadPoolExecutor(max_workers=2)
        # 用于在事件循环中安全地执行音频操作
        self.loop = None
    
    # 音频播放任务 - 异步优化版本
    async def play_audio(self):
        # 获取当前事件循环
        self.loop = asyncio.get_event_loop()
        
        p = None
        stream = None
        consecutive_errors = 0
        max_consecutive_errors = 3
        last_audio_time = time.time()
        audio_buffer = []  # 添加缓冲区以平滑播放
        buffer_target_size = 5  # 增加到5个数据包，提供更稳定的播放
        buffer_initialized = False  # 标记缓冲区是否已初始化
        
        try:
            while True:
                try:
                    # 初始化PyAudio
                    if p is None or stream is None:
                        # 使用线程池初始化PyAudio，避免阻塞事件循环
                        p, stream = await self.loop.run_in_executor(
                            self.executor, 
                            self._init_audio_stream
                        )
                        if p and stream:
                            self.is_playing = True
                            consecutive_errors = 0
                            audio_buffer = []  # 重置缓冲区
                        else:
                            consecutive_errors += 1
                            if consecutive_errors >= max_consecutive_errors:
                                print(f"[ERROR] 连续 {max_consecutive_errors} 次无法打开音频流，将在5秒后重试")
                                await asyncio.sleep(5)
                                consecutive_errors = 0
                            continue
                    
                    # 尝试从异步队列获取音频数据
                    try:
                        # 使用较短的超时以保持响应性
                        audio_data = await asyncio.wait_for(self.audio_queue.get(), timeout=0.5)
                        last_audio_time = time.time()
                        
                        # 验证数据完整性
                        if audio_data and len(audio_data) > 0:
                            # 将数据添加到缓冲区
                            audio_buffer.append(audio_data)
                            
                            # 当缓冲区达到目标大小时开始播放
                            if len(audio_buffer) >= buffer_target_size:
                                if not buffer_initialized:
                                    print(f"[AUDIO] 缓冲区已满({len(audio_buffer)}个包)，开始播放")
                                    buffer_initialized = True
                                
                                # 播放缓冲区中的数据
                                for buffered_data in audio_buffer:
                                    try:
                                        # 确保数据长度正确
                                        expected_size = CHUNK_SIZE * 2  # 每个样本2字节
                                        if len(buffered_data) != expected_size:
                                            # 调整数据长度
                                            if len(buffered_data) < expected_size:
                                                # 填充
                                                buffered_data = buffered_data + b'\x00' * (expected_size - len(buffered_data))
                                            else:
                                                # 截断
                                                buffered_data = buffered_data[:expected_size]
                                        
                                        # 使用线程池执行阻塞的写入操作
                                        await self.loop.run_in_executor(
                                            self.executor, 
                                            stream.write, 
                                            buffered_data
                                        )
                                    except Exception as write_e:
                                        print(f"[ERROR] 写入音频流失败: {write_e}")
                                        consecutive_errors += 1
                                        # 重新初始化音频设备
                                        if stream is not None:
                                            await self.loop.run_in_executor(self.executor, self._close_stream, stream)
                                            stream = None
                                        audio_buffer = []  # 清空缓冲区
                                        buffer_initialized = False
                                        break
                                
                                # 清空缓冲区
                                audio_buffer = []
                        else:
                            print(f"[WARN] 收到空或无效的音频数据")
                        
                        # 标记任务完成
                        self.audio_queue.task_done()
                    
                    except asyncio.TimeoutError:
                        # 队列为空，检查是否有缓冲数据需要播放
                        if audio_buffer and len(audio_buffer) > 0:
                            # 播放缓冲区中的剩余数据
                            print(f"[AUDIO] 队列空，播放剩余的 {len(audio_buffer)} 个缓冲包")
                            for buffered_data in audio_buffer:
                                try:
                                    # 确保数据长度正确
                                    expected_size = CHUNK_SIZE * 2
                                    if len(buffered_data) != expected_size:
                                        if len(buffered_data) < expected_size:
                                            buffered_data = buffered_data + b'\x00' * (expected_size - len(buffered_data))
                                        else:
                                            buffered_data = buffered_data[:expected_size]
                                    
                                    await self.loop.run_in_executor(
                                        self.executor, 
                                        stream.write, 
                                        buffered_data
                                    )
                                except:
                                    pass
                            audio_buffer = []
                            buffer_initialized = False
                        
                        # 检查是否需要关闭空闲的音频流
                        if time.time() - last_audio_time > 30 and self.is_playing:
                            print("[AUDIO] 长时间无音频数据，关闭空闲音频流")
                            if stream is not None:
                                await self.loop.run_in_executor(self.executor, self._close_stream, stream)
                                stream = None
                            self.is_playing = False
                    except Exception as queue_e:
                        print(f"[ERROR] 从音频队列获取数据失败: {queue_e}")
                    
                    # 每次循环后让出控制权，即使没有错误
                    await asyncio.sleep(0)
                    consecutive_errors = 0
                    
                except Exception as e:
                    print(f"[ERROR] 音频播放任务异常: {e}")
                    consecutive_errors += 1
                    
                    # 重新初始化音频设备
                    if stream is not None:
                        await self.loop.run_in_executor(self.executor, self._close_stream, stream)
                        stream = None
                    audio_buffer = []  # 清空缓冲区
                    buffer_initialized = False
                    
                    # 错误处理和重试逻辑
                    if consecutive_errors >= max_consecutive_errors:
                        print(f"[ERROR] 连续 {max_consecutive_errors} 次音频播放错误，将在5秒后重试")
                        await asyncio.sleep(5)
                        consecutive_errors = 0
                    else:
                        await asyncio.sleep(0.1)  # 短暂等待后重试
        
        finally:
            # 清理资源
            if stream is not None:
                try:
                    await self.loop.run_in_executor(self.executor, self._close_stream, stream)
                except Exception as e:
                    print(f"[ERROR] 关闭音频流失败: {e}")
            
            if p is not None:
                try:
                    p.terminate()
                except Exception as e:
                    print(f"[ERROR] 终止PyAudio失败: {e}")
            
            # 关闭线程池
            self.executor.shutdown(wait=False)
            print("[AUDIO] 音频播放任务已清理资源并退出")
    
    # 初始化音频流的辅助方法 - 在线程池中执行
    def _init_audio_stream(self):
        p = None
        stream = None
        try:
            p = pyaudio.PyAudio()
            
            # 获取默认输出设备
            try:
                default_device_index = int(p.get_default_output_device_info()['index'])  # type: ignore
                device_info = p.get_device_info_by_index(default_device_index)
                print(f"[AUDIO] 使用默认输出设备: {device_info['name']}")
            except Exception as device_e:
                print(f"[WARN] 无法获取默认设备信息: {device_e}")
            
            # 打开音频流
            stream = p.open(
                format=DEFAULT_FORMAT,
                channels=DEFAULT_CHANNELS,
                rate=DEFAULT_RATE,
                output=True,
                frames_per_buffer=CHUNK_SIZE
            )
            print(f"[AUDIO] 音频播放流已打开 - {DEFAULT_RATE}Hz, {DEFAULT_CHANNELS}声道")
            return p, stream
        except Exception as stream_e:
            print(f"[ERROR] 无法打开音频流: {stream_e}")
            if stream is not None:
                try:
                    self._close_stream(stream)
                except:
                    pass
            if p is not None:
                try:
                    p.terminate()
                except:
                    pass
            return None, None
    
    # 关闭音频流的辅助方法 - 在线程池中执行
    def _close_stream(self, stream):
        try:
            stream.stop_stream()
            stream.close()
        except Exception as e:
            print(f"[ERROR] 关闭音频流时出错: {e}")
    
    # 播放录制的音频
    def play_recorded_audio_with_data(self, recorded_data):
        print("[AUDIO] 开始播放录制的音频...")
        
        # 检查是否有录制的数据
        if not recorded_data:
            print("[AUDIO] 没有录制的音频数据")
            return
        
        print(f"[AUDIO] 准备播放 {len(recorded_data)} 个音频块")
        
        try:
            # 创建独立的PyAudio实例用于播放
            p = pyaudio.PyAudio()
            stream = None
            
            # 打开音频输出流
            stream = p.open(
                format=DEFAULT_FORMAT,
                channels=DEFAULT_CHANNELS,
                rate=DEFAULT_RATE,
                output=True,
                frames_per_buffer=CHUNK_SIZE
            )
            
            # 播放所有录制的音频数据
            for i, data in enumerate(recorded_data):
                stream.write(data)
                # 每10个块显示一次进度
                if (i + 1) % 10 == 0:
                    print(f"[AUDIO] 已播放 {i + 1}/{len(recorded_data)} 个音频块")
            
            print("[AUDIO] 录制音频播放完成")
            
        except Exception as e:
            print(f"[ERROR] 播放录制的音频时出错: {e}")
        finally:
            # 确保关闭流和PyAudio实例
            if stream is not None:
                stream.stop_stream()
                stream.close()
            if p is not None:
                p.terminate()
    
    # 测试音频播放功能
    def test_audio_playback(self):
        print("\n[TEST] 开始音频播放测试...")
        
        try:
            # 创建PyAudio实例
            p = pyaudio.PyAudio()
            
            # 打开音频流
            stream = p.open(
                format=DEFAULT_FORMAT,
                channels=DEFAULT_CHANNELS,
                rate=DEFAULT_RATE,
                output=True,
                frames_per_buffer=CHUNK_SIZE
            )
            
            print(f"[OK] 音频测试流已打开 - {DEFAULT_RATE}Hz, {DEFAULT_CHANNELS}声道")
            
            # 生成一个简单的正弦波
            duration = 0.5  # 0.5秒的测试音
            frequency = 440  # A4音
            samples = int(DEFAULT_RATE * duration)
            t = np.linspace(0, duration, samples, endpoint=False)
            
            # 创建正弦波并转换为16位PCM
            sine_wave = np.sin(2 * np.pi * frequency * t)
            # 调整音量到10%
            sine_wave = sine_wave * 0.1
            # 转换为16位整数
            pcm_data = np.int16(sine_wave * 32767)
            
            # 播放测试音
            print("[AUDIO] 播放测试音频 (440Hz 正弦波)")
            stream.write(pcm_data.tobytes())
            
            # 关闭流
            stream.stop_stream()
            stream.close()
            p.terminate()
            
            print("[OK] 音频测试成功完成！")
            return True
        except Exception as e:
            print(f"[ERROR] 音频测试失败: {e}")
            import traceback
            traceback.print_exc()
            return False

# asyncio已在文件顶部导入